/*
 * benchmark.cpp
 *
 * Copyright 2013 Long Tran (ltran3@gatech.edu)
 *
 * main file for maxtrix factorization program
 * Solving 0.5|A-WH^T|^2 + <Lambda, W-U> + <Pi, H-V> + rho/2 (|W-U|^2 + |H-V|^2)
 *  Generate data
 *   mpirun -np 16 ./nmf  --gen --data=m1000_10_200_16 --m=200 --n=200 --k=10 
 *  Sequential NMF (ALS)
 *   mpirun -np 1 ./nmf  --anls --m=200 --n=200 --k=10 --iter=1000 --rho=10
 *  Master/Slave NMF
 *   mpirun -np 4 ./nmf --anls-dist --k=10 --data=m1000_10_200_16 --rho=14 --iter=5000 --interval=1000
 *  Decentralized NMF (specify network)
 *   mpirun -np 4 ./nmf --anls-dist --k=10 --data=m1000_10_200_16 --rho=14 --graph=grid --iter=5000 --interval=1000
 *   mpirun -np 4 ./nmf --anls-dist --k=10 --data=m1000_10_200_16 --rho=14 --graph=equal --degree=2 --iter=5000 --interval=1000
 * 
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
 * MA 02110-1301, USA.
 * 
 */

#include <armadillo>
#include "helper/helper.h"
#include <elemental.hpp>
#include "algorithm/StopWatch.h"

using namespace elem;
typedef double R;
typedef elem::Complex<R> C;

po::variables_map vm;
po::options_description desc("Allowed options");

#define SECONDS(d) ((double) d.total_milliseconds() / 1000)

void process_program_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce help message")
    ("verbose", "print all debug information")
    ("interval", po::value<int>()->default_value(100), "data dimension hint")
    ("nodist", "only distributed Elemental")
    // DATA
    ("data", po::value<std::string>(), "data file")
    ("output", po::value<std::string>(), "output basis file")
    ("iterinfo", po::value<std::string>(), "iteration information file")
    // MODEL
    ("model", po::value<std::string>(), "output model file")
    // METHODS
    ("gen", "generate random data (MPI_SIZE)\n--m --n --k")
    ("svd", "benchmark SVD")
    ("mul", "benchmark Multiplication")
    ("solve", "benchmark Solve")
    ("chol", "benchmark Cholesky Solve")
    // COMMON PARAMETERS
    ("iter", po::value<int>()->default_value(1000), "maximum iterations")
    ("rate", po::value<double>()->default_value(1.0), "learning rate / step size")
    ("density", po::value<double>()->default_value(0.3), "sparsity density")
    ("rho", po::value<double>()->default_value(0.1), "augmented penalty")
    ("alpha", po::value<double>()->default_value(0.1), "concensus penalty")
    ("normalize", "keep columns of H normalized (norm <= 1)")
    ("l1", po::value<double>(), "l1 regularization")
    ("l2", po::value<double>(), "l2 regularization")
    // DISTRIBUTED ALGORITHMS
    ("graph", po::value<std::string>(), "decentralized ADMM with bipartie graph (grid|equal)")
    ("degree", po::value<int>(), "decentralized ADMM with bipartie graph")
    // GENERATE DATA
    ("m", po::value<int>()->default_value(10), "= A.nrows")
    ("n", po::value<int>()->default_value(10), "= A.ncols")
    ("k", po::value<int>()->default_value(2), "= W.ncols = H.nrows")
    ("np", po::value<int>()->default_value(2), "Number of processes to generate data")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm); 
}

int commRank = -1, commSize = -1;

class Benchmark {
  arma::mat& A_arma;
  elem::DistMatrix<R>& A_dist;
  elem::Matrix<R>& A_local;
  elem::Grid& g;
  StopWatch watch;
public:
  Benchmark(arma::mat& A_arma_, elem::DistMatrix<R>& A_dist_, elem::Matrix<R>& A_local_, elem::Grid& g_)
    : A_arma(A_arma_), A_dist(A_dist_), A_local(A_local_), g(g_)
  {
  }

  void SVD()
  {
    if (vm.count("nodist")) {
      watch.GetDuration();
      arma::vec sArma = arma::svd(A_arma);
      if (commRank==0) MSG("ARMA SVD d=" << SECONDS(watch.GetDuration()));
    }

    DistMatrix<R,VR,STAR> sOnly( g );
    watch.GetDuration();
    SingularValues( A_dist, sOnly );
    if (commRank==0) MSG("ELEM SVD d=" << SECONDS(watch.GetDuration()) << " (" << commSize << " core(s))");

    if (vm.count("nodist")) {
      Matrix<R> sOnly_local;
      watch.GetDuration();
      SingularValues( A_local, sOnly_local );
      if (commRank==0) MSG("ELEM local SVD d=" << SECONDS(watch.GetDuration()));
    }
  }

  void Mul()
  {
    if (vm.count("nodist")) {
      arma::mat B;
      watch.GetDuration();
      B = arma::trans(A_arma)*A_arma;
      if (commRank==0) MSG("ARMA Mul d=" << SECONDS(watch.GetDuration()));
    }

    DistMatrix<R> B(A_dist);
    Zero(B);
    watch.GetDuration();
    Gemm(TRANSPOSE, NORMAL, R(1), A_dist, A_dist, R(1), B);
    if (commRank==0) MSG("ELEM Mul d=" << SECONDS(watch.GetDuration())  << " (" << commSize << " core(s))");

    if (vm.count("nodist")) {
      Matrix<R> B(A_local);
      Zero(B);
      watch.GetDuration();
      Gemm(TRANSPOSE, NORMAL, R(1), A_local, A_local, R(1), B);
      if (commRank==0) MSG("ELEM local Mul d=" << SECONDS(watch.GetDuration()));
    }
  }

  void Solve()
  {
    if (vm.count("nodist")) {;
      watch.GetDuration();
      arma::mat X = arma::solve(A_arma, A_arma);
      if (commRank==0) MSG("ARMA Solve d=" << SECONDS(watch.GetDuration()));

      //Matrix<R> b_local(A_arma.n_rows, 1);
      //watch.GetDuration();
      //GaussianElimination(A_local, b_local);
      //if (commRank==0) MSG("ELEM local Solve (Gauss) d=" << SECONDS(watch.GetDuration()));
    }
    
    DistMatrix<R> b_dist(A_dist);
    watch.GetDuration();
    GaussianElimination(A_dist, b_dist);
    if (commRank==0) MSG("ELEM Distributed Solve (Gauss) d=" << SECONDS(watch.GetDuration())  << " (" << commSize << " core(s))");
  }

  void CholSolve()
  {
    if (vm.count("nodist")) {
      arma::mat X;
      watch.GetDuration();
      arma::mat U = chol(A_arma);
      arma::mat Y = arma::solve(arma::trimatl(arma::trans(U)), A_arma);
      X = arma::solve(arma::trimatu(U), Y);
      if (commRank==0) MSG("ARMA Solve d=" << SECONDS(watch.GetDuration()));

      Matrix<R> b_local( A_local );
      watch.GetDuration();
      CholeskySolve(UPPER, A_local, b_local);
      if (commRank==0) MSG("ELEM local Solve (Gauss) d=" << SECONDS(watch.GetDuration()));
    }
    
    DistMatrix<R> b_dist( A_dist );
    watch.GetDuration();
    CholeskySolve(UPPER, A_dist, b_dist);
    if (commRank==0) MSG("ELEM Distributed Solve (Gauss) d=" << SECONDS(watch.GetDuration())  << " (" << commSize << " core(s))");
  }

};

void gendata()
{
  PO_VARIABLE(m, vm, "m", int);
  //  PO_VARIABLE(n, vm, "n", int);
  
  arma::mat H = arma::randn<arma::mat>(m,m);
  arma::mat A = arma::trans(H)*H + 5*arma::eye(m,m);

  PO_VARIABLE(dataFile, vm, "data", std::string);
  A.save(dataFile, arma::arma_binary);
}

void readMatrix(arma::mat& A_arma, elem::DistMatrix<R>& A_dist, elem::Matrix<R>& A_local)
{
  PO_VARIABLE(data, vm, "data", std::string);
  A_arma.load(data, arma::arma_binary);
  int m = A_arma.n_rows, n = A_arma.n_cols;
  
  Uniform(m, n, A_dist);
  Uniform(m, n, A_local);
  for (int i=0; i < m; i++)
    for (int j=0; j < n; j++) {
      A_dist.Set(i,j, A_arma(i,j));
      A_local.Set(i,j, A_arma(i,j));
    }
}

int main(int argc, char** argv)
{
  Initialize( argc, argv );
  mpi::Comm comm = mpi::COMM_WORLD;
  commRank = mpi::CommRank( comm );
  commSize = mpi::CommSize( comm );

  process_program_options(argc, argv);

  try {
    const int nb = 96;
    SetBlocksize( nb );
    Grid g( comm );
    arma::mat A_arma;
    DistMatrix<R> A_dist(g);
    Matrix<R> A_local;

    if (vm.count("gen")) {
      gendata();
    }
    else {
      readMatrix(A_arma, A_dist, A_local);
      Benchmark b(A_arma, A_dist, A_local, g);

      if (vm.count("svd")) b.SVD();
      if (vm.count("mul")) b.Mul();
      if (vm.count("solve")) b.Solve();
      if (vm.count("chol")) b.CholSolve();
    }
  }
  catch ( std::exception& e ) {
    std::cerr << "Rank=" << commRank << " error: " << e.what() << std::endl;
  }
  Finalize();  
  return 0;
}
